package com.demo.util;


import com.demo.utils.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.driver.ShardingSphereDriverURL;
import org.apache.shardingsphere.infra.util.yaml.YamlEngine;
import org.apache.shardingsphere.infra.yaml.config.pojo.YamlRootConfiguration;
import org.springframework.util.Assert;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * 分片算法工具类
 */
@Slf4j
public class ShardingAlgorithmUtil {
    // 表分片符号。如：tb_order_202407
    private static final String TABLE_SPLIT_SYMBOL = "_";

    // 配置的数据库源
    private volatile static Map<String, Map<String, Object>> dataSources = null;

    public static void init(String url) {
        Assert.hasText(url, "分片策略不能为空");
        log.info("数据源获取...");
        byte[] bytes = new ShardingSphereDriverURL(url).toConfigurationBytes();
        try {
            YamlRootConfiguration yamlRootConfiguration = YamlEngine.unmarshal(bytes, YamlRootConfiguration.class);
            dataSources = yamlRootConfiguration.getDataSources();
        } catch(Exception e) {
            e.printStackTrace();
            log.error("分片策略配置解析失败");
            throw new IllegalArgumentException("分片策略解析失败");
        }
    }

    /**
     * 获取所有真实表名
     */
    public static Set<String> getAllTableNameBySchema(String logicTableName) {
        Assert.notNull(dataSources, "分片策略配置未初始化");
        Set<String> rs = new HashSet<>();
        // 获取配置的数据源
        String startTable = logicTableName + TABLE_SPLIT_SYMBOL;
        for (Map<String, Object> dataSource : dataSources.values()) {
            try (Connection conn = DriverManager.getConnection(dataSource.get("url").toString(),
                    dataSource.get("username").toString(), dataSource.get("password").toString())){
                Statement statement = conn.createStatement();
                ResultSet resultSet = statement.executeQuery("show tables like '" + startTable + "%'");
                while (resultSet.next()) {
                    String tableName = resultSet.getString(1);
                    if(StringUtils.hasText(tableName) && tableName.replaceFirst(startTable, "").matches("\\d{6}")) {
                        rs.add(tableName);
                    }
                }
            } catch(Exception e) {
                e.printStackTrace();
                throw new IllegalArgumentException("数据库连接失败");
            }
        }
        return rs;
    }

    /**
     * 创建分表
     * @param logicTableName
     * @param resultTableName
     * @return
     */
    public static boolean createShardingTable(String logicTableName, String resultTableName) {
        synchronized (logicTableName.intern()) {
            for (Map<String, Object> dataSource : dataSources.values()) {
                try (Connection conn = DriverManager.getConnection(dataSource.get("url").toString(),
                        dataSource.get("username").toString(), dataSource.get("password").toString())){
                    Statement statement = conn.createStatement();
                    log.info("创建{}表", resultTableName);
                    statement.execute("create table if not exists `" + resultTableName + "` like `" + logicTableName + "`;");
                } catch(Exception e) {
                    e.printStackTrace();
                    throw new IllegalArgumentException("数据库连接失败");
                }
            }
            return true;
        }
    }

}
